from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

x = pd.read_excel("E:\本科教学\大三下\数学建模\新疆大学数学建模校赛2021年度赛题及说明\A题\data.xlsx", header=None)
x = np.array(x)
y = np.array(x[:, 5]).reshape(len(x), 1)
x = np.vstack((np.vstack((x[:, 1], x[:, 2])))).transpose()
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=2)

reg = LinearRegression()
reg.fit(x_train,y_train,)
y_test_predict = reg.predict(x_test)

print("截距：",reg.coef_)
print("权重：",reg.intercept_)
print("均方误差：",mean_squared_error(y_test, y_test_predict))

fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(x[:,0],x[:,1],y,s=10,c="black")

ax.scatter(x[:,0], x[:,1], y, s=10, c="black")
xx = np.arange(0, 300, 0.5)
yy = np.arange(0, 250, 0.5)
X, Y = np.meshgrid(xx, yy)
Z = reg.intercept_[0] + reg.coef_.reshape(2)[0] * X + reg.coef_.reshape(2)[1] * Y
ax.plot_surface(X, Y, Z, cmap='rainbow')
ax.set_zlabel('Z', fontdict={'size': 15, 'color': 'red'})
ax.set_ylabel('Y', fontdict={'size': 15, 'color': 'red'})
ax.set_xlabel('X', fontdict={'size': 15, 'color': 'red'})

plt.show()

